﻿//
// C# KD Tree Implementation from https://code.google.com/p/kd-sharp/
// Based on the Java implementation from : https://bitbucket.org/rednaxela/knn-benchmark/src/tip/ags/utils/dataStructures/trees/thirdGenKD/ </remarks>
//
using System;
using System.Collections;
using System.Collections.Generic;
using System.Linq;
using System.Text;

namespace KDTreeRednaxela
{
    /// <summary>
    /// A NearestNeighbour iterator for the KD-tree which intelligently iterates and captures relevant data in the search space.
    /// </summary>
    /// <typeparam name="T">The type of data the iterator should handle.</typeparam>
    public class NearestNeighbour<T> : IEnumerator
    {
        /// <summary>The point from which are searching in n-dimensional space.</summary>
        private float[] tSearchPoint;
        /// <summary>A distance function which is used to compare nodes and value positions.</summary>
        private IDistanceFunction kDistanceFunction;
        /// <summary>The tree nodes which have yet to be evaluated.</summary>
        private MinHeap<KDNode_Rednaxela<T>> pPending;
        /// <summary>The values which have been evaluated and selected.</summary>
        private IntervalHeap<T> pEvaluated;
        /// <summary>The root of the kd tree to begin searching from.</summary>
        private KDNode_Rednaxela<T> pRoot = null;

        /// <summary>The max number of points we can return through this iterator.</summary>
        private int iMaxPointsReturned = 0;
        /// <summary>The number of points we can still test before conclusion.</summary>
        private int iPointsRemaining;
        /// <summary>Threshold to apply to tree iteration.  Negative numbers mean no threshold applied.</summary>
        private float fThreshold;

        /// <summary>Current value distance.</summary>
        private float _CurrentDistance = -1;
        /// <summary>Current value reference.</summary>
        private T _Current = default(T);
        /// <summary>Current index distance.</summary>
        private int _CurrentIndex = -1;
        /// <summary>
        /// Construct a new nearest neighbour iterator.
        /// </summary>
        /// <param name="pRoot">The root of the tree to begin searching from.</param>
        /// <param name="tSearchPoint">The point in n-dimensional space to search.</param>
        /// <param name="kDistance">The distance function used to evaluate the points.</param>
        /// <param name="iMaxPoints">The max number of points which can be returned by this iterator.  Capped to max in tree.</param>
        /// <param name="fThreshold">Threshold to apply to the search space.  Negative numbers indicate that no threshold is applied.</param>
        public NearestNeighbour(KDNode_Rednaxela<T> pRoot, float[] tSearchPoint, IDistanceFunction kDistance, int iMaxPoints, float fThreshold)
        {
            // Check the dimensionality of the search point.
            if (tSearchPoint.Length != pRoot.dimensions)
                throw new Exception("Dimensionality of search point and kd-tree are not the same.");

            // Store the search point.
            this.tSearchPoint = new float[tSearchPoint.Length];
            Array.Copy(tSearchPoint, this.tSearchPoint, tSearchPoint.Length);

            // Store the point count, distance function and tree root.
            this.iPointsRemaining = Math.Min(iMaxPoints, pRoot.Size);
            this.fThreshold = fThreshold;
            this.kDistanceFunction = kDistance;
            this.pRoot = pRoot;
            this.iMaxPointsReturned = iMaxPoints;
            _CurrentDistance = -1;

            // Create an interval heap for the points we check.
            this.pEvaluated = new IntervalHeap<T>();

            // Create a min heap for the things we need to check.
            this.pPending = new MinHeap<KDNode_Rednaxela<T>>();
            this.pPending.Insert(0, pRoot);
        }
        private void CheckSinglePoint(KDNode_Rednaxela<T> pCursor)
        {
            // Work out the distance between this point and the search point.
            float fDistance = kDistanceFunction.Distance(pCursor.points[0], tSearchPoint);

            //// Skip if the point exceeds the threshold.
            //// Technically this should never happen, but be prescise.
            //if (fThreshold >= 0 && fDistance >= fThreshold)
            //    continue;

            // Add the point if either need more points or it's closer than furthest on list so far.
            if (pEvaluated.Size < iPointsRemaining || fDistance <= pEvaluated.MaxKey)
            {
                for (int i = 0; i < pCursor.Size; ++i)
                {
                    // If we don't need any more, replace max
                    if (pEvaluated.Size == iPointsRemaining)
                        pEvaluated.ReplaceMax(fDistance, pCursor.data[i], i);

                    // Otherwise insert.
                    else
                        pEvaluated.Insert(fDistance, pCursor.data[i], i);
                }
            }
        }
        private void CheckSpreadOutPoints(KDNode_Rednaxela<T> pCursor)
        {
            // Treat the distance of each point seperately.
            for (int i = 0; i < pCursor.Size; ++i)
            {
                // Compute the distance between the points.
                float fDistance = kDistanceFunction.Distance(pCursor.points[i], tSearchPoint);

                //// Skip if it exceeds the threshold.
                //if (fThreshold >= 0 && fDistance >= fThreshold)
                //    continue;

                // Insert the point if we have more to take.
                if (pEvaluated.Size < iPointsRemaining)
                    pEvaluated.Insert(fDistance, pCursor.data[i], i);

                // Otherwise replace the max.
                else if (fDistance < pEvaluated.MaxKey)
                    pEvaluated.ReplaceMax(fDistance, pCursor.data[i], i);
            }
        }
        private void DescendPath(KDNode_Rednaxela<T> pCursor)
        {
            KDNode_Rednaxela<T> pNotTaken;

            // If the seach point is larger, select the right path.
            if (tSearchPoint[pCursor.splitDimension] > pCursor.fSplitValue)
            {
                pNotTaken = pCursor.pLeft;
                pCursor = pCursor.pRight;
            }
            else
            {
                pNotTaken = pCursor.pRight;
                pCursor = pCursor.pLeft;
            }

            // Calculate the shortest distance between the search point and the min and max bounds of the kd-node.
            float fDistance = kDistanceFunction.DistanceToRectangle(tSearchPoint, pNotTaken.minBound, pNotTaken.maxBound);

            //// If it is greater than the threshold, skip.
            //if (fThreshold >= 0 && fDistance > fThreshold)
            //{
            //    //pPending.Insert(fDistance, pNotTaken);
            //    continue;
            //}

            // Only add the path we need more points or the node is closer than furthest point on list so far.
            if (pEvaluated.Size < iPointsRemaining || fDistance <= pEvaluated.MaxKey)
            {
                pPending.Insert(fDistance, pNotTaken);
            }
        }
        /// <summary>
        /// Check for the next iterator item.
        /// </summary>
        /// <returns>True if we have one, false if not.</returns>
        public bool MoveNext()
        {
            // Bail if we are finished.
            if (iPointsRemaining == 0)
            {
                _Current = default(T);
                return false;
            }

            // While we still have paths to evaluate.
            while (pPending.Size > 0 && (pEvaluated.Size == 0 || (pPending.MinKey < pEvaluated.DistanceMin)))
            {
                // If there are pending paths possibly closer than the nearest evaluated point, check it out
                KDNode_Rednaxela<T> pCursor = pPending.Min;
                pPending.RemoveMin();

                // Descend the tree, record paths not taken
                while (!pCursor.IsLeaf)
                {
                   // DescendPath(pCursor);
                    KDNode_Rednaxela<T> pNotTaken;

                    // If the seach point is larger, select the right path.
                    if (tSearchPoint[pCursor.splitDimension] > pCursor.fSplitValue)
                    {
                        pNotTaken = pCursor.pLeft;
                        pCursor = pCursor.pRight;
                    }
                    else
                    {
                        pNotTaken = pCursor.pRight;
                        pCursor = pCursor.pLeft;
                    }

                    // Calculate the shortest distance between the search point and the min and max bounds of the kd-node.
                    float fDistance = kDistanceFunction.DistanceToRectangle(tSearchPoint, pNotTaken.minBound, pNotTaken.maxBound);

                    //float fDistance = kDistanceFunction.Distance(tSearchPoint, pNotTaken.maxBound);


                    //// If it is greater than the threshold, skip.
                    //if (fThreshold >= 0 && fDistance > fThreshold)
                    //{
                    //    //pPending.Insert(fDistance, pNotTaken);
                    //    continue;
                    //}

                    // Only add the path we need more points or the node is closer than furthest point on list so far.
                    if (pEvaluated.Size < iPointsRemaining || fDistance <= pEvaluated.MaxKey)
                    {
                        pPending.Insert(fDistance, pNotTaken);
                    }
                }

                // If all the points in this KD node are in one place.
                if (pCursor.IsSinglePoint)
                {
                    CheckSinglePoint(pCursor);
                }

                // If the points in the KD node are spread out.
                else
                {
                    CheckSpreadOutPoints(pCursor);
                }
            }

            // Select the point with the smallest distance.
            if (pEvaluated.Size == 0)
                return false;

            iPointsRemaining--;
            _CurrentDistance = pEvaluated.DistanceMin;
            _Current = pEvaluated.ClosestPoint;
            _CurrentIndex = pEvaluated.ClosestPointIndex;
            pEvaluated.RemoveMin();
            return true;
        }

        /// <summary>
        /// Reset the iterator.
        /// </summary>
        public void Reset()
        {
            // Store the point count and the distance function.
            this.iPointsRemaining = Math.Min(iMaxPointsReturned, pRoot.Size);
            _CurrentDistance = -1;

            // Create an interval heap for the points we check.
            this.pEvaluated = new IntervalHeap<T>();

            // Create a min heap for the things we need to check.
            this.pPending = new MinHeap<KDNode_Rednaxela<T>>();
            this.pPending.Insert(0, pRoot);
        }

        /// <summary>
        /// Return the current value referenced by the iterator as an object.
        /// </summary>
        object IEnumerator.Current
        {
            get { return _Current; }
        }

        /// <summary>
        /// Return the distance of the current value to the search point.
        /// </summary>
        public float CurrentDistance
        {
            get { return _CurrentDistance; }
        }

        /// <summary>
        /// Return the current value referenced by the iterator.
        /// </summary>
        public T CurrentPoint
        {
            get { return _Current; }
        }
        /// <summary>
        /// Return the current value referenced by the iterator.
        /// </summary>
        public int CurrentIndex
        {
            get { return _CurrentIndex; }
        }
    }
}
